
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.model_selection import train_test_split
import pandas as pd
from sklearn import naive_bayes as nb
from sklearn.metrics import accuracy_score,recall_score,precision_score,f1_score

# 1 读取数据

df = pd.read_csv('C:\\Users\\12948\\Desktop\\垃圾邮件分类\\train.csv', engine='python')

df['encoded_label']=df.Label.map({'spam':0,'ham':1})#把原先的ham（正常邮件）和spam（垃圾邮件）用1和0代替


# 2 随机划分训练集train_data
train_data, test_data, train_label, test_label = train_test_split(
    df.Email,
    df.encoded_label,
    test_size=0.3,#设置测试集占30%
    random_state=0)  # df.Email是邮件内容，df.encoded_label是邮件标签（ham和spam）


# 3 文本特征抽取，使用CountVectorizer将句子转化为向量
c_v = CountVectorizer(decode_error='ignore')
train_data = c_v.fit_transform(train_data)
test_data = c_v.transform(test_data)


# 4 朴素贝叶斯算法训练预测
clf=nb.MultinomialNB()#初始化模型

model=clf.fit(train_data, train_label)#输入训练集对模型进行训练

predicted_label=model.predict(test_data)#输入测试集进行分类测试


# 5 评估模型
print('准确率:', format(accuracy_score(test_label, predicted_label)))
print('精确率:', format(precision_score(test_label, predicted_label)))
print('召回率:', format(recall_score(test_label, predicted_label)))
print('F1 分数:', format(f1_score(test_label, predicted_label)))
